# -*- coding: utf-8 -*-

import psycopg2
import yaml
from typing import List, Dict, Any, Optional, Tuple


class PostgresConnector:
    def __init__(self, config_path: str = 'config.yml'):
        self.config = self._load_config(config_path)
        self.connection = None
        self.cursor = None
        
    def _load_config(self, config_path: str) -> Dict[str, Any]:
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
        return config['postgres']
    
    def connect(self) -> None:
        if self.connection is None or self.connection.closed:
            self.connection = psycopg2.connect(
                host=self.config['addr'],
                port=self.config['port'],
                dbname=self.config['db-name'],
                user=self.config['username'],
                password=self.config['password']
            )
            self.cursor = self.connection.cursor()
    
    def disconnect(self) -> None:
        if self.cursor:
            self.cursor.close()
        if self.connection and not self.connection.closed:
            self.connection.close()
        self.cursor = None
        self.connection = None
    
    def execute_query(self, query: str, params: Optional[Tuple] = None) -> List:
        self.connect()
        self.cursor.execute(query, params or ())
        try:
            result = self.cursor.fetchall()
            return result
        except psycopg2.ProgrammingError:
            return []
    
    def list_databases(self) -> List[str]:
        query = "SELECT datname FROM pg_database WHERE datistemplate = false;"
        result = self.execute_query(query)
        return [row[0] for row in result]
    
    def list_tables(self) -> List[str]:
        query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';"
        result = self.execute_query(query)
        return [row[0] for row in result]
    
    def list_columns(self, table_name: str) -> List[Dict[str, str]]:
        query = """
        SELECT column_name, data_type 
        FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = %s;
        """
        result = self.execute_query(query, (table_name,))
        return [{'column': row[0], 'type': row[1]} for row in result]
    
    def list_all_tables_columns(self) -> Dict[str, List[Dict[str, str]]]:
        tables = self.list_tables()
        result = {}
        for table in tables:
            result[table] = self.list_columns(table)
        return result
    
    def get_table_content(self, table_name: str, limit: int = 10) -> List[Dict]:
        if limit <= 0:
            limit = 10
            
        query = f"""
        SELECT * FROM {table_name} 
        ORDER BY (SELECT 1) 
        LIMIT %s;
        """
        columns = self.list_columns(table_name)
        column_names = [col['column'] for col in columns]
        
        result = self.execute_query(query, (limit,))
        return [dict(zip(column_names, row)) for row in result]
    
    def get_table_first_last_records(self, table_name: str, count: int = 10) -> Dict[str, List[Dict]]:
        columns = self.list_columns(table_name)
        column_names = [col['column'] for col in columns]
        
        first_query = f"""
        SELECT * FROM {table_name} 
        ORDER BY (SELECT 1) 
        LIMIT %s;
        """
        
        last_query = f"""
        SELECT * FROM {table_name} 
        ORDER BY (SELECT 1) DESC 
        LIMIT %s;
        """
        
        first_result = self.execute_query(first_query, (count,))
        last_result = self.execute_query(last_query, (count,))
        
        return {
            'first': [dict(zip(column_names, row)) for row in first_result],
            'last': [dict(zip(column_names, row)) for row in last_result]
        }
    
    def __enter__(self):
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect() 